Jax ml operator fix by Ig-dolci · Pull Request #4041 · firedrakeproject/firedrake

Ig-dolci · 2025-02-14T11:51:59Z

Description

Fixed an issue in firedrake.ml.jax.ml_operator where argument_slots was incorrectly specified. I wrote it as optional:
```
argument_slots: Optional[tuple[Union[ufl.coefficient.BaseCoefficient, ufl.argument.BaseArgument]]] = ()
```
Clarified that if argument_slots is not provided, ML_Operator will automatically write it.
Test the results involving the Neural operators are in right function space.
Test jax and pytorch operators in Firedrake CI.

github-actions · 2025-02-14T12:13:26Z

	Tests	Passed ✅	Skipped ⏭️	Failed ❌
Firedrake complex	8269 ran	6595 passed	1674 skipped	0 failed

github-actions · 2025-02-14T12:13:40Z

	Tests	Passed ✅	Skipped ⏭️	Failed ❌
Firedrake real	8198 ran	7500 passed	698 skipped	0 failed

firedrake/ml/jax/ml_operator.py

…orch and jax.

firedrake/ml/pytorch/ml_operator.py

firedrake/external_operators/ml_operator.py

firedrake/ml/pytorch/ml_operator.py

firedrake/external_operators/abstract_external_operators.py

firedrake/ml/pytorch/ml_operator.py

firedrake/external_operators/ml_operator.py

firedrake/ml/jax/ml_operator.py

firedrake/ml/pytorch/ml_operator.py

scripts/firedrake-install

Ig-dolci · 2025-02-21T10:33:13Z

@pbrubeck @dham Firedrake pip installs PR depends on the approval of UFL PR 348 and this PR.

Ig-dolci · 2025-02-21T10:37:17Z

The failure in the docs building is because of a recent Sphinx update.

tests/firedrake/external_operators/test_external_operators_adjoint.py

firedrake/ml/jax/ml_operator.py

.github/workflows/build.yml

firedrake/ml/jax/fem_operator.py

connorjward

@pbrubeck if you are also happy it would be fantastic if this could be merged ASAP as it is holding up the pip PR. Naturally we need to update our UFL fork at the same time.

.github/workflows/build.yml

Does this should be optional?

dc14436

Ig-dolci changed the title ~~Does this should be optional?~~ Jax ml operator fix Feb 14, 2025

A hack before understand argument_slots

c838418

pbrubeck reviewed Feb 15, 2025

View reviewed changes

firedrake/ml/jax/ml_operator.py Outdated Show resolved Hide resolved

Ig-dolci added 2 commits February 17, 2025 13:59

Checking argument_slots; Adding docs; testing actions results for pyt…

e6d87f6

…orch and jax.

More docs

ee001ce

Ig-dolci marked this pull request as ready for review February 17, 2025 14:11

Fix docs

7278949

pbrubeck reviewed Feb 17, 2025

View reviewed changes

firedrake/ml/pytorch/ml_operator.py Outdated Show resolved Hide resolved

dham reviewed Feb 17, 2025

View reviewed changes

firedrake/external_operators/ml_operator.py Outdated Show resolved Hide resolved

Ig-dolci added 3 commits February 18, 2025 13:26

Enable CI tests for optional installed packages

8249bf4

docs

96fc8e5

docs

e48501c

pbrubeck reviewed Feb 18, 2025

View reviewed changes

firedrake/ml/pytorch/ml_operator.py Outdated Show resolved Hide resolved

Ig-dolci added 14 commits February 19, 2025 12:30

Test ufl Pablo's fixing

09003f6

merge master

df1ee6b

More assert; Increase tolerance to reduce test runtime

3aa5939

just a test

1469403

install jaxlib

aca9860

Test another jax install requirement for linux

ef6e4d2

wip

d52431c

wip

a50cea8

wip

a094d88

wip

22be7c4

será?

6d3b726

será?

f11a99e

test

b5f7080

test

24c757c

Ig-dolci added 3 commits February 20, 2025 15:13

Use options={riesz_representation: None}

fe649de

Jax install fixing

65289a2

Merge branch 'master' into dolci/jax_ml_operator

2c4f227

Ig-dolci commented Feb 20, 2025

View reviewed changes

Update scripts/firedrake-install

fa4c380

Ig-dolci commented Feb 21, 2025

View reviewed changes

tests/firedrake/external_operators/test_external_operators_adjoint.py Show resolved Hide resolved

connorjward reviewed Feb 21, 2025

View reviewed changes

firedrake/ml/jax/ml_operator.py Outdated Show resolved Hide resolved

pbrubeck reviewed Feb 21, 2025

View reviewed changes

.github/workflows/build.yml Outdated Show resolved Hide resolved

Switch ufl branch

8b50efd

pbrubeck reviewed Feb 21, 2025

View reviewed changes

firedrake/ml/jax/fem_operator.py Show resolved Hide resolved

Ig-dolci added 3 commits February 21, 2025 17:19

review

0e00379

review

6c477d8

review

d613eb2

connorjward previously approved these changes Feb 25, 2025

View reviewed changes

pbrubeck reviewed Feb 25, 2025

View reviewed changes

.github/workflows/build.yml Outdated Show resolved Hide resolved

Update .github/workflows/build.yml

6762c85

pbrubeck dismissed connorjward’s stale review via 6762c85 February 25, 2025 13:13

pbrubeck approved these changes Feb 25, 2025

View reviewed changes

pbrubeck enabled auto-merge (squash) February 25, 2025 13:13

Ig-dolci disabled auto-merge February 25, 2025 13:58

Ig-dolci enabled auto-merge (squash) February 25, 2025 14:00

connorjward disabled auto-merge February 25, 2025 15:14

connorjward merged commit a84c071 into master Feb 25, 2025
19 of 20 checks passed

connorjward deleted the dolci/jax_ml_operator branch February 25, 2025 15:15

Conversation

Ig-dolci commented Feb 14, 2025 • edited by pbrubeck Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Uh oh!

github-actions bot commented Feb 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Feb 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Ig-dolci commented Feb 21, 2025

Uh oh!

Ig-dolci commented Feb 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

connorjward left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Ig-dolci commented Feb 14, 2025 •

edited by pbrubeck

Loading

github-actions bot commented Feb 14, 2025 •

edited

Loading

github-actions bot commented Feb 14, 2025 •

edited

Loading

Ig-dolci commented Feb 21, 2025 •

edited

Loading